#include "integrate5.h"
#include "common.h"
#include "common_cuda.h"
#include <pthread.h>

#define BLOCK_WIDTH 512

typedef struct DataStruct {
        unsigned int deviceID;
        unsigned int deviceCount;
    bodyList* bl;
    unsigned int iterations;
    unsigned int restNum;
    unsigned int GRID_WIDTH;
} DataStruct;

//barrier for sychronizing the threads on each iteration
pthread_barrier_t mybarrier;

//__device__ tile
__device__ float3
tile_calculation5(float4* myPosition, float4* shPosition, float3* accel, unsigned int tileID, unsigned int restNum)
{
    int num = blockDim.x;

    if(restNum != 0 && tileID == (gridDim.x-1))
        num = restNum;

    int i = 0;
    for(i = 0; i < num; i++){
        *accel += calcGravityAccel(myPosition, &shPosition[i]);
    }

    return *accel;
}

//__global__ do the calculation
__global__ void 
updateVel5(float4* globalX, float4* globalA, int N, unsigned int restNum, unsigned int deviceID, unsigned int deviceCount)
{
    __shared__ float4 shPosition[BLOCK_WIDTH];

    int i = 0;
    int tile = 0;
    int gtid = blockIdx.x*blockDim.x + threadIdx.x;
    float3 acc = make_float3(0.0f);
    float4 myPosition = globalX[gtid];
    
    int offset = N/deviceCount;
    int lowerRange = offset*deviceID;
    int upperRange = lowerRange + offset;
    
    bool didCal = false;

    for(i = 0, tile = 0; i < N; i +=blockDim.x, tile++){
        
        if(gtid < N){
            int idx = tile*blockDim.x + threadIdx.x;
            if(idx < N){
            shPosition[threadIdx.x] = globalX[idx];
            }
        }


        __syncthreads();
        
        //do the partitioned calculations based on the range
        if((lowerRange <= gtid) && (gtid < upperRange)){
            acc = tile_calculation5(&myPosition, shPosition, &acc, tile, restNum);
            didCal = true;
        }
        
        __syncthreads();
    }
        //save the result in global memory for the integration step
    if(didCal){
        globalA[gtid] += make_float4(acc);
    }
}

void * ptr_func(void* dataStruct){

    DataStruct* data= (DataStruct*)dataStruct;
    unsigned int deviceID = data->deviceID;
    unsigned int n = data->bl->size;
    size_t size = n*sizeof(float4);
    
    //set current device
    cudaSetDevice(deviceID);
    
    float4* d_pos;
    float4* d_vel;
    
    // allocate memory on gpu and copy data
    d_pos = (float4*) safe_cudaMalloc(size);
    d_vel = (float4*) safe_cudaMalloc(size);
    cudaMemcpy(d_pos, data->bl->pos, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_vel, data->bl->vel, size, cudaMemcpyHostToDevice);
    
    //determine the grid and block size
    dim3 dimb(BLOCK_WIDTH, 1);
    dim3 dimg(data->GRID_WIDTH, 1);

    //invoke kernel
    for (int i = 0; i != data->iterations; ++i)
    {
        updateVel5<<<dimg, dimb>>>(d_pos, d_vel, n, data->restNum, deviceID, data->deviceCount);
        
        cudaThreadSynchronize();
        pthread_barrier_wait(&mybarrier);
        
        updatePos<<<dimg, dimb>>>(d_pos, d_vel, n, 1.0f);
        
        cudaThreadSynchronize();
        pthread_barrier_wait(&mybarrier);
    }

    // copy results from gpu back to cpu
    unsigned int partial_size = n/(data->deviceCount);
    cudaMemcpy(data->bl->pos + (deviceID*partial_size), d_pos + (deviceID*partial_size), partial_size*sizeof(float4), cudaMemcpyDeviceToHost);
    cudaMemcpy(data->bl->vel + (deviceID*partial_size), d_vel + (deviceID*partial_size), partial_size*sizeof(float4), cudaMemcpyDeviceToHost);

    // free gpu memory
    cudaFree(d_pos);
    cudaFree(d_vel);
    
    return 0;
}

void integrate5(bodyList* bl, int iterations){
    int deviceCount = 0;
    cudaGetDeviceCount(&deviceCount);
    
    //calculate grid width and see if there needs an extra block
    unsigned int GRID_WIDTH = (bl->size)/BLOCK_WIDTH;
    unsigned int restNum = (bl->size)%BLOCK_WIDTH;
    if(restNum!=0)
        GRID_WIDTH += 1;
        
    //set up data struct
    DataStruct dataStruct[deviceCount];
    
    int i = 0;
    for(i = 0; i<deviceCount; i++ ){
        dataStruct[i].deviceID = i;
        dataStruct[i].deviceCount = deviceCount;
        dataStruct[i].bl = bl;
        dataStruct[i].iterations = iterations;
        dataStruct[i].restNum = restNum;
        dataStruct[i].GRID_WIDTH = GRID_WIDTH;
    }

    //run pthreads
    pthread_t thread[deviceCount];
    //initiate my barrier
    pthread_barrier_init(&mybarrier, NULL, deviceCount);
    
    i = 0;
    for(i = 0; i< deviceCount; i++){
        pthread_create(&thread[i], NULL, ptr_func, (void*)&(dataStruct[i]) );
    }
    
    //join pthreads
    i = 0;
    for(i = 0; i< deviceCount; i++){
        pthread_join(thread[i], NULL);
    }
    //destroy barrier
    pthread_barrier_destroy(&mybarrier);

}
